from torchvision.models.resnet import BasicBlock
import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np

from model.model_baseline import Net
from utils import to_var, softmax_cross_entropy_criterion

BatchNorm2d = nn.BatchNorm2d


class FusionNet(nn.Module):
    def load_pretrain(self, pretrain_file):
        # raise NotImplementedError
        pretrain_state_dict = torch.load(pretrain_file)
        state_dict = self.state_dict()
        keys = list(state_dict.keys())
        for key in keys:
            state_dict[key] = pretrain_state_dict[key]

        self.load_state_dict(state_dict)
        print('')

    def __init__(self, num_class=2):
        super(FusionNet, self).__init__()

        self.color_moudle = Net(num_class=num_class, is_first_bn=True)

        self.depth_moudle = Net(num_class=num_class, is_first_bn=True)

        self.ir_moudle = Net(num_class=num_class, is_first_bn=True)

        self.res_0 = self._make_layer(BasicBlock, 384, 256, 2, stride=2)
        self.res_1 = self._make_layer(BasicBlock, 256, 512, 2, stride=2)

        self.fc = nn.Sequential(nn.Dropout(0.5),
                                nn.Linear(512, 256),
                                nn.ReLU(inplace=True),
                                nn.Linear(256, num_class))

    def _make_layer(self, block, inplanes, planes, blocks, stride=1):
        downsample = None
        if stride != 1:
            downsample = nn.Sequential(
                nn.Conv2d(inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion), )

        layers = []
        layers.append(block(inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        batch_size, c, h, w = x.shape

        color = x[:, 0:3, :, :]
        depth = x[:, 3:6, :, :]
        ir = x[:, 6:9, :, :]

        color_feas = self.color_moudle.forward_res3(color)
        depth_feas = self.depth_moudle.forward_res3(depth)
        ir_feas = self.ir_moudle.forward_res3(ir)
        fea = torch.cat([color_feas, depth_feas, ir_feas], dim=1)

        x = self.res_0(fea)
        x = self.res_1(x)
        x = F.adaptive_avg_pool2d(x, output_size=1).view(batch_size, -1)
        x = self.fc(x)
        return x, None, None

    def set_mode(self, mode, is_freeze_bn=False):
        self.mode = mode
        if mode in ['eval', 'valid', 'test']:
            self.eval()
        elif mode in ['backup']:
            self.train()
            if is_freeze_bn:  # freeze
                for m in self.modules():
                    if isinstance(m, BatchNorm2d):
                        m.eval()
                        m.weight.requires_grad = False
                        m.bias.requires_grad = False


def run_check_net():
    batch_size = 32
    c, h, w = 3, 128, 128
    num_class = 2

    nn_input = np.random.uniform(0, 1, (batch_size, c, h, w)).astype(np.float32)
    nn_truth = np.random.choice(num_class, batch_size).astype(np.float32)

    # ------------
    nn_input = torch.from_numpy(nn_input).float().cuda()
    nn_truth = torch.from_numpy(nn_truth).long().cuda()

    nn_input = to_var(nn_input)
    nn_truth = to_var(nn_truth)

    # ---
    criterion = softmax_cross_entropy_criterion
    net = Net(num_class).cuda()
    net.set_mode('backup')
    print(net)

    logit = net.forward(nn_input)
    loss = criterion(logit, nn_truth)


if __name__ == '__main__':
    import os

    os.environ['CUDA_VISIBLE_DEVICES'] = '4,5,6,7'  # '3,2,1,0'
    print('%s: calling main function ... ' % os.path.basename(__file__))
    run_check_net()
    print('sucessful!')
